package com.intelligent.system.auth.uitls;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

import javax.annotation.PostConstruct;
import java.math.BigInteger;
import java.security.KeyFactory;
import java.security.PublicKey;
import java.security.spec.RSAPublicKeySpec;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;

/**
 * OAuth2证书工具类
 * 用于获取和验证OAuth2服务器的公钥证书
 */
@Slf4j
@Component
public class OAuth2CertUtils {

    @Value("${spring.security.oauth2.client.registration.my-auth-server.client-id}")
    private String clientId;

    @Value("${spring.security.oauth2.client.registration.my-auth-server.client-secret}")
    private String clientSecret;

    @Value("${spring.security.oauth2.client.provider.my-auth-server.jwk-set-uri}")
    private String certUri;

    private final RestTemplate restTemplate = new RestTemplate();
    private final ObjectMapper objectMapper = new ObjectMapper();
    private PublicKey publicKey;
    private String keyId;

    @PostConstruct
    public void init() {
        try {
            refreshPublicKey();
        } catch (Exception e) {
            log.error("初始化公钥失败", e);
        }
    }

    /**
     * 刷新公钥
     * 从OAuth2服务器获取最新的公钥证书
     */
    public void refreshPublicKey() {
        try {
            // 构建请求头
            HttpHeaders headers = new HttpHeaders();
            headers.set("Content-Type", "application/json");

            // 构建请求参数
            Map<String, String> params = new HashMap<>();
            headers.set("client_id", clientId);
            headers.set("client_secret", clientSecret);

            // 发送请求获取证书
            HttpEntity<Map<String, String>> requestEntity = new HttpEntity<>(params, headers);
            ResponseEntity<String> response = restTemplate.exchange(
                    certUri,
                HttpMethod.POST,
                requestEntity,
                String.class
            );

            // 解析响应
            JsonNode root = objectMapper.readTree(response.getBody());
            if (root.has("code") && root.get("code").asInt() == 200) {
                JsonNode keys = root.get("data").get("keys");
                if (keys.isArray() && keys.size() > 0) {
                    JsonNode key = keys.get(0);
                    this.keyId = key.get("kid").asText();
                    this.publicKey = parseJwkPublicKey(key);
                    log.info("成功刷新公钥，密钥ID: {}", this.keyId);
                } else {
                    throw new RuntimeException("No public key found in response");
                }
            } else {
                log.error("获取公钥失败: {}", root.get("msg").asText());
                throw new RuntimeException("Failed to get public key");
            }
        } catch (Exception e) {
            log.error("刷新公钥时发生错误", e);
            throw new RuntimeException("Failed to refresh public key", e);
        }
    }

    /**
     * 获取当前公钥
     */
    public PublicKey getPublicKey() {
        if (publicKey == null) {
            refreshPublicKey();
        }
        return publicKey;
    }

    /**
     * 获取当前密钥ID
     */
    public String getKeyId() {
        if (keyId == null) {
            refreshPublicKey();
        }
        return keyId;
    }

    /**
     * 解析JWK格式的公钥
     * JWK格式包含RSA公钥的模数(n)和指数(e)
     */
    private PublicKey parseJwkPublicKey(JsonNode key) throws Exception {
        // 获取RSA公钥的模数(n)和指数(e)
        String n = key.get("n").asText();
        String e = key.get("e").asText();

        // 解码Base64URL编码的模数和指数
        byte[] nBytes = Base64.getUrlDecoder().decode(n);
        byte[] eBytes = Base64.getUrlDecoder().decode(e);

        // 转换为BigInteger
        BigInteger modulus = new BigInteger(1, nBytes);
        BigInteger exponent = new BigInteger(1, eBytes);

        // 创建RSA公钥规范
        RSAPublicKeySpec keySpec = new RSAPublicKeySpec(modulus, exponent);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        return keyFactory.generatePublic(keySpec);
    }
} 